/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lvyh.lightframe.tracer.plugin.kafka;

import com.lvyh.lightframe.tracer.core.SimpleSpan;
import com.lvyh.lightframe.tracer.core.SimpleTracer;
import com.lvyh.lightframe.tracer.core.common.AbstractTracer;
import com.lvyh.lightframe.tracer.core.common.SimpleTracerConfiguration;
import com.lvyh.lightframe.tracer.core.common.TraceContants;
import com.lvyh.lightframe.tracer.core.context.SimpleTraceContext;
import com.lvyh.lightframe.tracer.core.context.SimpleTraceContextHolder;
import com.lvyh.lightframe.tracer.core.samplers.Sampler;
import com.lvyh.lightframe.tracer.plugin.kafka.carrier.KafkaMqInjectCarrier;
import io.opentracing.propagation.Format;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.clients.producer.Callback;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.apache.kafka.common.Metric;
import org.apache.kafka.common.MetricName;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.ProducerFencedException;
import org.apache.kafka.common.header.Headers;

import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

public class TracerKafkaProducer<K, V> implements Producer<K, V> {

    private Producer<K, V> producer;

    private AbstractTracer kafkaSendTracer;

    public TracerKafkaProducer(Producer<K, V> producer, AbstractTracer kafkaSendTracer) {
        this.producer = producer;
        this.kafkaSendTracer = kafkaSendTracer;
    }

    @Override
    public void initTransactions() {
        producer.initTransactions();
    }

    @Override
    public void beginTransaction() throws ProducerFencedException {
        producer.beginTransaction();
    }

    @Override
    public void sendOffsetsToTransaction(Map<TopicPartition, OffsetAndMetadata> offsets, String consumerGroupId) throws ProducerFencedException {
        producer.sendOffsetsToTransaction(offsets, consumerGroupId);
    }

    @Override
    public void commitTransaction() throws ProducerFencedException {
        producer.commitTransaction();
    }

    @Override
    public void abortTransaction() throws ProducerFencedException {
        producer.abortTransaction();
    }

    @Override
    public Future<RecordMetadata> send(ProducerRecord<K, V> producerRecord) {
        return send(producerRecord, null);
    }

    @Override
    public Future<RecordMetadata> send(ProducerRecord<K, V> producerRecord, Callback callback) {
        SimpleSpan clientSpan = kafkaSendTracer.clientSend("mq-kafka-produce-msg");

        setSpanTags(producerRecord, clientSpan);

        injectCarrier(clientSpan, producerRecord.headers());

        SimpleTraceContext simpleTraceContext = SimpleTraceContextHolder.getSimpleTraceContext();
        if (clientSpan.getParentSpan() != null) {
            simpleTraceContext.add(clientSpan.getParentSpan());
        }

        return producer.send(producerRecord, new TracerCallback(callback, kafkaSendTracer, clientSpan));
    }

    @Override
    public void flush() {
        producer.flush();
    }

    @Override
    public List<PartitionInfo> partitionsFor(String topic) {
        return producer.partitionsFor(topic);
    }

    @Override
    public Map<MetricName, ? extends Metric> metrics() {
        return producer.metrics();
    }

    @Override
    public void close() {
        producer.close();
    }

    @Override
    public void close(long timeout, TimeUnit timeUnit) {
        producer.close(timeout, timeUnit);
    }

    @Override
    public void close(Duration duration) {
    }

    private void setSpanTags(ProducerRecord<K, V> producerRecord, SimpleSpan clientSpan) {
        clientSpan.setTag(TraceContants.CURRENT_THREAD_NAME, Thread.currentThread().getName());
        clientSpan.setTag(TraceContants.LOCAL_APP, SimpleTracerConfiguration.getApplicationName()+ "#Producer");

        clientSpan.setTag(TraceContants.KAFKA_TOPIC, producerRecord.topic());
        clientSpan.setTag(TraceContants.KAFKA_PARTITION, producerRecord.partition() == null ? -1 : producerRecord.partition());
    }

    private void injectCarrier(SimpleSpan simpleSpan, Headers properties) {
        SimpleTracer simpleTracer = this.kafkaSendTracer.getSimpleTracer();
        simpleTracer.inject(simpleSpan.getSimpleSpanContext(), Format.Builtin.TEXT_MAP, new KafkaMqInjectCarrier(properties));
    }

    static final class TracerCallback implements Callback {

        final Callback callback;

        final AbstractTracer kafkaSendTracer;

        private SimpleSpan simpleSpan;

        public TracerCallback(Callback callback, AbstractTracer kafkaSendTracer) {
            this.callback = callback;
            this.kafkaSendTracer = kafkaSendTracer;
        }

        public TracerCallback(Callback callback, AbstractTracer kafkaSendTracer, SimpleSpan simpleSpan) {
            this.simpleSpan = simpleSpan;
            this.callback = callback;
            this.kafkaSendTracer = kafkaSendTracer;
        }

        @Override
        public void onCompletion(RecordMetadata metadata, Exception exception) {
            boolean successFlag = true;
            if (Objects.nonNull(exception)) {
                successFlag = false;
            }
            if (Objects.nonNull(callback)) {
                callback.onCompletion(metadata, exception);
            }

            if (simpleSpan != null) {
                SimpleSpan currentSpan = simpleSpan;
                Sampler sampler = kafkaSendTracer.getSimpleTracer().getSampler();
                //sampler is support &  current span is root span
                if (sampler != null && currentSpan.getParentSpan() == null) {
                    currentSpan.getSimpleSpanContext().setSampled(true);
                    currentSpan.getSimpleSpanContext().setSampler(sampler);
                }
            }

            kafkaSendTracer.clientReceiveTagFinish(this.simpleSpan,
                    successFlag ? TraceContants.RESULT_CODE_SUCCESS : TraceContants.RESULT_CODE_ERROR);
        }
    }
}

